package com.iwdnb.gkgz.common.quota;

import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import com.google.common.collect.Lists;
import com.iwdnb.gkgz.common.model.dto.StockDayData;
import com.iwdnb.gkgz.common.utils.BigDecimalUtils;
import com.iwdnb.gkgz.common.utils.StockUtils;

/**
 * 实现股票的 RSI（相对强弱指数）指标，首先需要定义一个代表股票日K线数据的类，然后能根据该类中的收盘价计算某一段时间内的 RSI 值。
 *
 * 下面我们具体步骤如下：
 *
 * 定义一个代表股票日K线数据的类 StockDayInfo。
 * 创建一个方法，计算一段时间内的 RSI 值。
 */
public class RSIIndicator {
    public static final String QUOTA_NAME = "rsi";
    private static final Integer PERIOD = 6;

    public static void calculateRSI(List<StockDayData> stockDayDataList) {
        List<Double> list = doCalculateRSI(stockDayDataList, PERIOD);
        int size = stockDayDataList.size();
        for (int i = 0; i < size; i++) {
            StockDayData stockDayData = stockDayDataList.get(i);
            Double d = list.get(i);
            BigDecimal data = BigDecimalUtils.of(d);
            StockUtils.setStockQuota(stockDayData, QUOTA_NAME, data);
        }
    }

    private static List<Double> doCalculateRSI(List<StockDayData> stockDataList, int period) {
        List<Double> rsiList = new ArrayList<>();
        if (stockDataList == null || stockDataList.size() < period) {
            return rsiList; // 缺乏足够的数据来计算RSI
        }

        double[] closes = new double[stockDataList.size()];
        for (int i = 0; i < stockDataList.size(); i++) {
            closes[i] = stockDataList.get(i).getClosePrice().doubleValue();
        }

        double gainSum = 0;
        double lossSum = 0;

        // 初始化第一个周期的涨幅和跌幅
        for (int i = 1; i <= period; i++) {
            double change = closes[i] - closes[i - 1];
            if (change > 0) {
                gainSum += change;
            } else {
                lossSum += Math.abs(change);
            }
        }

        double avgGain = gainSum / period;
        double avgLoss = lossSum / period;

        for (int i = period; i < closes.length; i++) {
            double change = closes[i] - closes[i - 1];
            if (change > 0) {
                avgGain = (avgGain * (period - 1) + change) / period;
                avgLoss = (avgLoss * (period - 1)) / period;
            } else {
                avgGain = (avgGain * (period - 1)) / period;
                avgLoss = (avgLoss * (period - 1) + Math.abs(change)) / period;
            }

            double rs = avgGain / avgLoss;
            double rsi = 100 - (100 / (1 + rs));

            rsiList.add(rsi);
        }

        // 前面的period个时间点没有RSI值，因此使用0填充
        for (int i = 0; i < period; i++) {
            rsiList.add(0, 0d);
        }

        return rsiList;
    }

    public static void main(String[] args) {
        // 示例股票数据
        List<StockDayData> stockDataList = Lists.newArrayList(
            createStockDayData("2023-10-01", 100.0),
            createStockDayData("2023-10-02", 102.0),
            createStockDayData("2023-10-03", 104.0),
            createStockDayData("2023-10-04", 103.0),
            createStockDayData("2023-10-05", 101.0),
            createStockDayData("2023-10-06", 99.0),
            createStockDayData("2023-10-07", 98.0),
            createStockDayData("2023-10-08", 97.0),
            createStockDayData("2023-10-09", 96.0),
            createStockDayData("2023-10-10", 95.0),
            createStockDayData("2023-10-01", 96),
            createStockDayData("2023-10-02", 94),
            createStockDayData("2023-10-03", 98),
            createStockDayData("2023-10-04", 103.0),
            createStockDayData("2023-10-05", 107),
            createStockDayData("2023-10-06", 108),
            createStockDayData("2023-10-07", 110),
            createStockDayData("2023-10-08", 105),
            createStockDayData("2023-10-09", 100),
            createStockDayData("2023-10-10", 98),
            createStockDayData("2023-10-10", 95.0),
            createStockDayData("2023-10-01", 86),
            createStockDayData("2023-10-02", 84),
            createStockDayData("2023-10-03", 78),
            createStockDayData("2023-10-04", 73.0),
            createStockDayData("2023-10-05", 67),
            createStockDayData("2023-10-06", 68),
            createStockDayData("2023-10-07", 63),
            createStockDayData("2023-10-08", 59),
            createStockDayData("2023-10-09", 56),
            createStockDayData("2023-10-10", 52)
        );

        // 计算RSI，14天的周期是常见的选择
        int period = 14;
        List<Double> rsiList = doCalculateRSI(stockDataList, period);

        // 输出结果
        for (int i = 0; i < rsiList.size(); i++) {
            System.out.println("RSI: " + rsiList.get(i));
        }
    }

    private static StockDayData createStockDayData(String date, double price) {
        StockDayData stockDayInfo = new StockDayData();
        stockDayInfo.setClosePrice(new BigDecimal(price));
        return stockDayInfo;
    }
}

